package utils

import (
	"encoding/base64"
	"fmt"
	"os"
	"reflect"
	"strconv"
	"strings"
)

// 生成云主机自定义配置信息
// 如何注入参考 https://support.huaweicloud.com/usermanual-ecs/zh-cn_topic_0032380449.html
const (
	User_Data_File_Path = "C:\\UEOfflineRendering\\agent\\config.yaml"
)

type Config struct {
	MySQL      *MySQLConfig  `yaml:"mysql"`
	Filevolume *[]FileVolume `yaml:"filevolumes"`
}

type MySQLConfig struct {
	Host     string `yaml:"host"`
	Port     int    `yaml:"port"`
	User     string `yaml:"user"`
	Password string `yaml:"passwd"`
	Database string `yaml:"db"`
}

type FileVolume struct {
	MountPath    string `yaml:"mount_path"`
	Addr         string `yaml:"addr"`
	MountCommand string `yaml:"args"`
}

func GenerateUserData(fileVolume *[]FileVolume) string {
	config := Config{
		MySQL: &MySQLConfig{
			Host:     os.Getenv("DB_HOST"),
			Port:     3306,
			User:     os.Getenv("DB_USER"),
			Password: os.Getenv("DB_PASSWORD"),
			Database: os.Getenv("DB_NAME_RENDERJOB"),
		},
		Filevolume: fileVolume,
	}

	portStr := os.Getenv("DB_PORT")
	port, err := strconv.Atoi(portStr)
	if err == nil {
		config.MySQL.Port = port
	}

	// 使用 strings.Builder 优化字符串拼接
	var builder strings.Builder

	// 添加脚本开头
	builder.WriteString("#ps1\n")

	// 将配置信息写入脚本
	writeConfig(&builder, "mysql", config.MySQL)
	writeConfigArray(&builder, "filevolumes", config.Filevolume)

	// 将脚本内容进行 Base64 编码
	return base64.StdEncoding.EncodeToString([]byte(builder.String()))
}

// 写入配置信息到脚本
func writeConfig(builder *strings.Builder, section string, data interface{}) {
	if data != nil {
		builder.WriteString(fmt.Sprintf("echo '%s:' >> %s\n", section, User_Data_File_Path))
		writeFields(builder, data)
	}
}

// 写入数组形式的配置信息到脚本
func writeConfigArray(builder *strings.Builder, section string, data *[]FileVolume) {
	if data != nil {
		builder.WriteString(fmt.Sprintf("echo '%s:' >> %s\n", section, User_Data_File_Path))
		for i, item := range *data {
			v := reflect.ValueOf(item)
			t := v.Type()
			for j := 0; j < v.NumField(); j++ {
				field := v.Field(j)
				tag := t.Field(j).Tag.Get("yaml")
				value := field.Interface()
				if j == 0 {
					builder.WriteString(fmt.Sprintf("echo '  - %s: \"%v\"' >> %s\n", tag, value, User_Data_File_Path))
				} else {
					builder.WriteString(fmt.Sprintf("echo '    %s: \"%v\"' >> %s\n", tag, value, User_Data_File_Path))
				}
			}
			if i < len(*data)-1 {
				builder.WriteString(fmt.Sprintf("echo '' >> %s\n", User_Data_File_Path))
			}
		}
	}
}

// 写入字段到脚本
func writeFields(builder *strings.Builder, data interface{}) {
	v := reflect.ValueOf(data)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}
	t := v.Type()
	for i := 0; i < v.NumField(); i++ {
		field := v.Field(i)
		tag := t.Field(i).Tag.Get("yaml")
		value := field.Interface()
		switch value.(type) {
		case int:
			builder.WriteString(fmt.Sprintf("echo '  %s: %v' >> %s\n", tag, value, User_Data_File_Path))
		default:
			builder.WriteString(fmt.Sprintf("echo '  %s: \"%v\"' >> %s\n", tag, value, User_Data_File_Path))
		}
	}
}

func GenerateUserDataForNative(dataVolumes, fileVolumes []string) string {
	var (
		initializeDisksFun string
		initializeDisks    string
		filevolumeFun      string
	)
	userData := getUserDataTmplForNative()
	if len(dataVolumes) > 0 {
		initializeDisksFun = `function InitializeDisks {
  Log "Get all RAW disks, initialize the disks, partition, and format them"
  Get-Disk | Where-Object PartitionStyle -eq "RAW" | Initialize-Disk -PassThru | New-Partition -AssignDriveLetter -UseMaximumSize | Format-Volume
  Log "Initialize RAW disks completed"
}`
		initializeDisks = "InitializeDisks"
	}
	if len(fileVolumes) > 0 {
		filevolumeFun = `  Log "Mouting lustre"`
		for i, dataVolume := range fileVolumes {
			diskTag := fmt.Sprintf("%c", 77+i)
			filevolumeFun += fmt.Sprintf(`
  launchctl-x64 start virtiofs viofs%s %s %s:	
`, diskTag, dataVolume, diskTag)
		}
	}

	return base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf(userData, initializeDisksFun, filevolumeFun, initializeDisks)))
}

func getUserDataTmplForNative() string {
	return `#ps1
param(
  [parameter()]
  [string]$MainPath,
  [string]$UserName,
  [string]$EncodedPaasword,
  [int]$Scheduler
)

$Password = $null
$AutoLogonPath="HKLM:\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Winlogon"
$AutoRunPath = "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Run"
$AutoRunKey = "UserDataCustomScript"

function Log([string]$Msg) {
  $LogFile = "$MainPath\Logs\UserDataCustomScriptLog.txt"
  $Msg = "$(Get-Date): Scheduler=$Scheduler, $Msg"
  Write-Host $Msg
  Add-Content -Path $LogFile -Value $Msg -Encoding Unicode
}

%s

function Startup {
  $OutFile = "$MainPath\Logs\UserDataCustomScriptOut.txt"
  Add-Content -Path $OutFile -Value "Startup...." -Encoding Unicode

  Log "Locking screen"
  . "$MainPath\Scripts\UserDataUtilsScript.ps1"
  [Utils]::LockWorkStation()
%s
  
}

function IsEnabledAutoLogon {
  $DefaultPassword = Get-ItemProperty -Path $AutoLogonPath -Name "DefaultPassword" -ErrorAction Ignore
  return $DefaultPassword -ne $null
}

function EnableAutoLogon {
  Log "EnableAutoLogon"
  Set-ItemProperty -Path $AutoLogonPath -Name "DefaultUserName" -Value "Administrator" -Force
  Set-ItemProperty -Path $AutoLogonPath -Name "DefaultPassword" -Value "$Password" -Force
  Set-ItemProperty -Path $AutoLogonPath -Name "AutoAdminLogon" -Value 1 -Force
}

function IsEnabledAutoRun {
  $AutoRunValue = Get-ItemProperty -Path $AutoRunPath -Name $AutoRunKey -ErrorAction Ignore
  return $AutoRunValue -ne $null
}

function EnableAutoRun {
  Log "EnableAutoRun"
  Set-ItemProperty -Path $AutoRunPath -Name $AutoRunKey -Value "PowerShell -ExecutionPolicy RemoteSigned -WindowStyle Hidden -File $MainPath\Scripts\UserDataCustomScript.ps1 -MainPath $MainPath -UserName $UserName -EncodedPaasword $EncodedPaasword" -Force
}

function EnableSchedulerTask {
  $Output = schtasks /Query | findstr /i "UserDataCustomScript"
  Log "Query schtasks output is '$Output'"
  if ($Output -eq $null) {
    Log "Create scheduler task"
    schtasks /Create /TN UserDataCustomScript /SC onstart /RU $UserName /RP $Password /TR "PowerShell -ExecutionPolicy RemoteSigned -WindowStyle Hidden -File $MainPath\Scripts\UserDataCustomScript.ps1 -MainPath $MainPath -UserName $UserName -EncodedPaasword $EncodedPaasword -Scheduler 1"
  }
}

function EnableNetExporter {
  $Output = schtasks /Query | findstr /i "NetExporter"
  if ($Output -eq $null) {
    schtasks /Create /TN NetExporter /SC onstart /RU $UserName /RP $Password /TR "C:\Program Files\net-exporter\net-exporter-windowsx86-64.exe"
  }
}

function InitParams {
  if ($MainPath -eq "") {
    $Script:MainPath = $Env:CloudUserData_MainPath
  }
  if ($UserName -eq "") {
    $Script:UserName = $Env:CloudUserData_UserName
  }
  if ($EncodedPaasword -eq "") {
    $Script:EncodedPaasword = $Env:CloudUserData_Password
  }
  Log "MainPath=$MainPath, UserName=$UserName, EncodedPaasword=$EncodedPaasword"
  $Script:Password = [Text.Encoding]::ASCII.GetString([Convert]::FromBase64String($EncodedPaasword))
}

function Reboot {
  Log "Reboot"
  shutdown /r /f /t 0
}

function EnableStartupAndReboot {
  EnableAutoLogon
  EnableAutoRun
  %s
  Reboot
}

function Main {
  InitParams
  Log "Starting UserDataCustomScript"
  try {
    if ($Scheduler -eq 1) {
      if (IsEnabledAutoLogon) {
        Log "Nothing"
      } else {
        EnableStartupAndReboot
      }
    } else {
      EnableSchedulerTask
      EnableNetExporter
      if (IsEnabledAutoRun) {
        Startup
      } else {
        EnableStartupAndReboot
      }
    }
  } catch {
    Log "An error occurred: $_"
    Log $_.ScriptStackTrace
  }
}
Main
`
}
